using Microsoft.EntityFrameworkCore;
using RAG.Domain.Entities.App;
using RAG.Domain.Repositories;
using RAG.Infrastructure.Data;

namespace RAG.Infrastructure.Repositories;

public class DocumentChunkRepository : IDocumentChunkRepository
{
    private readonly RagAIDbContext _context;

    public DocumentChunkRepository(RagAIDbContext context)
    {
        _context = context;
    }

    public async Task<DocumentChunks> CreateAsync(DocumentChunks chunk)
    {
        _context.DocumentChunks.Add(chunk);
        await _context.SaveChangesAsync();
        return chunk;
    }

    public async Task<DocumentChunks?> GetByIdAsync(Guid id)
    {
        return await _context.DocumentChunks
            .Include(c => c.Document)
            .Include(c => c.KnowledgeBase)
            .FirstOrDefaultAsync(c => c.Id == id);
    }

    public async Task<List<DocumentChunks>> GetByDocumentIdAsync(Guid documentId)
    {
        return await _context.DocumentChunks
            .Include(c => c.KnowledgeBase)
            .Where(c => c.Document.Id == documentId)
            .OrderBy(c => c.ChunkIndex)
            .ToListAsync();
    }

    public async Task<List<DocumentChunks>> GetChunksWithVectorsAsync(List<Guid> chunkIds)
    {
        return await _context.DocumentChunks
            .Include(c => c.Document)
            .Include(c => c.KnowledgeBase)
            .Where(c => chunkIds.Contains(c.Id))
            .ToListAsync();
    }

    public async Task DeleteByDocumentIdAsync(Guid documentId)
    {
        var chunks = await _context.DocumentChunks
            .Where(c => c.Document.Id == documentId)
            .ToListAsync();

        _context.DocumentChunks.RemoveRange(chunks);
        await _context.SaveChangesAsync();
    }
}
